# -*- coding: utf-8 -*-
"""
@Time ： 2023/8/24 18:52
@Auth ：gaoyu
@File ：output_data.py
@IDE ：PyCharm
"""

import traceback

import pandas as pd

from air_web.data_platform import mysql_con, sql_engine


def get_all_industry_data():
    sql = f"SELECT distinct (cons_no),shi,xian,type_code,type_id,pare_type_code,pare_type_id, run_cap FROM air_db_smxu_hebei.aclr_base_doc_all WHERE type_code_sort=1 and pare_type_id in (102,101,103,104,105,106,108,110,111,112,109) and c_id<100000000"
    result = mysql_con.get(sql)
    data = pd.DataFrame(result)

    data["第三级行业"] = data["pare_type_id"].apply(
        lambda x: "第一产业"
        if x in [102]
        else "第二产业"
        if x in [102, 109]
        else "第三产业"
    )
    data.rename(
        columns={
            "cons_no": "户号",
            "shi": "市区",
            "xian": "县区",
            "pare_type_code": "第四级行业",
            "run_cap": "运行容量",
        },
        inplace=True,
    )
    data = data[["户号", "市区", "县区", "第三级行业", "第四级行业", "运行容量"]]
    data["第二级行业"] = "全行业用电合计"
    data["第一级行业"] = "全社会用电总计"
    return data


def get_city_data():
    sql = f"SELECT distinct cons_no,shi,xian,type_code,type_id,pare_type_code,pare_type_id, run_cap FROM air_db_smxu_hebei.aclr_base_doc_all WHERE type_code_sort=1 and type_id in (10701,10702) and c_id <100000000"
    result = mysql_con.get(sql)
    data = pd.DataFrame(result)
    data["第三级行业"] = data["type_code"]
    data.rename(
        columns={
            "cons_no": "户号",
            "shi": "市区",
            "xian": "县区",
            "type_code": "第四级行业",
            "run_cap": "运行容量",
        },
        inplace=True,
    )
    data = data[["户号", "市区", "县区", "第三级行业", "第四级行业", "运行容量"]]
    data["第二级行业"] = "城乡居民用电合计"
    data["第一级行业"] = "全社会用电总计"
    return data


def main():
    try:
        data = get_all_industry_data()
        city_data = get_city_data()
        df = pd.concat([data, city_data])
        df["省"] = "河北省"
        df = df.reindex(
            columns=[
                "户号",
                "第一级行业",
                "第二级行业",
                "第三级行业",
                "第四级行业",
                "省",
                "市区",
                "县区",
                "运行容量",
            ]
        )
        print(df.columns)
        df.to_excel("result.xlsx", encoding="utf-8", index=False)
    except:
        print(traceback.format_exc())


if __name__ == "__main__":
    main()
